package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"log"
	"os"
	"strconv"
	"strings"
)

func getObj(results []entity.Column, name string) entity.Column {
	var obj entity.Column

	for _, result := range results {
		if result.Name() == name {
			obj = result
		}
	}
	return obj
}

// 辅助函数：将 float32 数组转换为逗号分隔的字符串
func float32ArrayToString(array []float32) string {
	var strArray []string
	for _, val := range array {
		strArray = append(strArray, strconv.FormatFloat(float64(val), 'f', -1, 32))
	}
	return strings.Join(strArray, ",")
}

func main() {
	// 创建 Milvus 客户端
	ctx := context.Background()
	//milvusClient, err := client.NewGrpcClient(ctx, "10.10.22.104:19530") // 修改为你的 Milvus 地址
	milvusClient, err := client.NewGrpcClient(ctx, "192.168.30.80:31530") // 修改为你的 Milvus 地址
	if err != nil {
		log.Fatalf("failed to create milvus client: %v", err)
	}
	defer milvusClient.Close()
	var limit = 10
	// 定义查询选项
	opt := client.SearchQueryOptionFunc(func(option *client.SearchQueryOption) {
		option.Limit = int64(limit)
		option.Offset = 20000
	})

	// 查询条件和需要输出的字段
	collectionName := "galleryNew"
	expr := "id > 0"
	outputFields := []string{
		"embedding", "image_type", "gallery_image_id", "path",
		"ai_model", "model_version", "prompt", "width", "height",
		"aspectRatio", "created_at", "deleted_at",
	}

	// 执行查询
	queryResults, err := milvusClient.Query(
		ctx, collectionName, nil, expr, outputFields, opt,
	)
	if err != nil {
		log.Fatalf("failed to query milvus: %v", err)
	}
	fmt.Println(queryResults)

	// 创建 CSV 文件
	file, err := os.Create("milvus_query_results.csv")
	if err != nil {
		log.Fatalf("failed to create CSV file: %v", err)
	}
	defer file.Close()

	// 初始化 CSV 写入器
	writer := csv.NewWriter(file)
	defer writer.Flush()

	// 写入表头
	writer.Write(outputFields)

	// 遍历查询结果并写入 CSV
	var embeddingObj = getObj(queryResults, "embedding")
	var imageTypeObj = getObj(queryResults, "image_type")
	var galleryImageIdObj = getObj(queryResults, "gallery_image_id")
	var pathObj = getObj(queryResults, "path")
	var aiModelObj = getObj(queryResults, "ai_model")
	var modelVersionObj = getObj(queryResults, "model_version")
	var promptObj = getObj(queryResults, "prompt")
	var widthObj = getObj(queryResults, "width")
	var heightObj = getObj(queryResults, "height")
	var aspectRatioObj = getObj(queryResults, "aspectRatio")
	var createdAtObj = getObj(queryResults, "created_at")
	var deletedAtObj = getObj(queryResults, "deleted_at")
	for i := 0; i < limit; i++ {

		// 假设 embedding 是一个 []float32 数组
		embedding, _ := embeddingObj.Get(i)
		embeddingArray, ok := embedding.([]float32)
		if !ok {
			log.Fatalf("unexpected type, expected []float32")
		}
		// 将 float32 切片转换为 string 切片
		strArray := make([]string, len(embeddingArray))
		for i, v := range embeddingArray {
			// 将 float32 转换为字符串，保留 6 位小数
			strArray[i] = strconv.FormatFloat(float64(v), 'f', 6, 32)
		}

		// 使用 strings.Join 并添加逗号分隔
		embeddingResult := "[" + strings.Join(strArray, ", ") + "]"

		imageType, _ := imageTypeObj.GetAsInt64(i)
		galleryImageId, _ := galleryImageIdObj.GetAsInt64(i)
		path, _ := pathObj.GetAsString(i)
		aiModel, _ := aiModelObj.GetAsString(i)
		modelVersion, _ := modelVersionObj.GetAsString(i)
		prompt, _ := promptObj.GetAsString(i)
		width, _ := widthObj.GetAsInt64(i)
		height, _ := heightObj.GetAsInt64(i)
		aspectRatio, _ := aspectRatioObj.GetAsDouble(i)
		createdAt, _ := createdAtObj.GetAsInt64(i)
		deletedAt, _ := deletedAtObj.GetAsInt64(i)
		record := []string{
			fmt.Sprintf("%v", embeddingResult),
			fmt.Sprintf("%v", imageType),
			fmt.Sprintf("%v", galleryImageId),
			fmt.Sprintf("%v", path),
			fmt.Sprintf("%v", aiModel),
			fmt.Sprintf("%v", modelVersion),
			fmt.Sprintf("%v", prompt),
			fmt.Sprintf("%v", width),
			fmt.Sprintf("%v", height),
			fmt.Sprintf("%v", aspectRatio),
			fmt.Sprintf("%v", createdAt),
			fmt.Sprintf("%v", deletedAt),
		}
		writer.Write(record)
	}

	fmt.Println("CSV 文件导出成功！")
}
